import pandas as pd
import numpy as np
import json
import io
# import time
# from factor_analyzer.factor_analyzer import calculate_kmo, calculate_bartlett_sphericity

def format_Change(originalData:list[dict], labelMap:dict =None) -> dict:
    NewData = {}
    for key in originalData[0].keys():
        key_data = []
        for data in originalData:
            key_data.append(data[key])
        if labelMap is not None and key in labelMap.keys():
            key = labelMap[key]
        if type(key_data[0]) is list:
            for i in range(0, len(key_data[0])):
                if i == len(key_data[0])-1 and key == '污染源':
                    key = '未识别源'
                    NewData[key] = [data[i] for data in key_data]
                    break
                NewData[key+'_'+str(i+1)] = [data[i] for data in key_data]
        else:
            NewData[key] = key_data
    return NewData

class PictureDrawer():
    def __init__(self, totalNum:int, ytitle=None, xlabel=None) -> None:
        self.linenums = 0
        self.totalNum = totalNum
        self.markers=["o", "v", "^", "<", ">", "1", "2", "3", "4", "8", "s", "p", "P", "*", "h", "H", "+", "x", "X", "D", "d"]
        self.colors=mpl.colormaps['tab20'].resampled(totalNum)
        self.ytitle = ytitle
        self.xlabel = xlabel
        self.x = []
        self.y = []
        self.labels = []
        plt.figure(figsize=(7,5))

    def addPlot(self, x, y, label):
        self.x.append(x)
        self.y.append(y)
        self.labels.append(label)

    def plot(self, x, y, label, marker=None, color=None, added=False) -> None:

        marker = self.markers if marker is None else marker
        color = self.colors if color is None else color

        if not added:
            self.addPlot(x, y, label)


        plt.plot(x, y, marker=marker[self.linenums], linestyle='', label=label, color=self.colors(self.linenums))
        self.linenums += 1

    def toBase64(self) -> str:

        for i in range(self.linenums, self.totalNum):
            self.plot(self.x[i], self.y[i], self.labels[i], added=True)

        lg = plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        frame = lg.get_frame()
        frame.set_alpha(1)
        frame.set_facecolor('none') # 设置图例legend背景透明

        plt.tight_layout()
        if self.ytitle is not None:
            plt.ylabel(self.ytitle)
        if self.xlabel is not None:
            plt.xticks(self.x[0], self.xlabel)

        buf = io.BytesIO()
        # plt.savefig(buf, format='png', dpi=600, bbox_extra_artists=(lg,), bbox_inches='tight', transparent=True)
        plt.savefig("test.png", format='png', dpi=600, bbox_extra_artists=(lg,), bbox_inches='tight', transparent=True)
        buf.seek(0)

        with open("test.png", "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
        return encoded_string

        # # 将字节流转换为Pillow格式
        # # im = Image.open(buf).convert('RGB')
        # im = Image.open(buf)

        # # 将Pillow格式的图像转换为Base64编码
        # im_buffer = io.BytesIO()
        # im.save(im_buffer, format='PNG')
        # im_bytes = im_buffer.getvalue()
        # return base64.b64encode(im_bytes).decode('utf-8')

    def clean(self):
        plt.clf()

class Region():
    def __init__(self, 
                 name, 
                 water_consumption, 
                 population, 
                 drainage_area, 
                 runoff_coefficient,
                 intercepting_multiplier,
                 minQuantityRainfall,
                 maxQuantityRainfall
                 ) -> None:
        self.name = name
        self.wc = water_consumption.astype('float64') if type(water_consumption) is np.str_ else water_consumption
        self.popu = population.astype('float64') if type(population) is np.str_ else population
        self.da = drainage_area.astype('float64') if type(drainage_area) is np.str_ else drainage_area
        self.rc = runoff_coefficient.astype('float64') if type(runoff_coefficient) is np.str_ else runoff_coefficient
        self.im = intercepting_multiplier.astype('float64') if type(intercepting_multiplier) is np.str_ else intercepting_multiplier
        self.minQR = minQuantityRainfall.astype('float64') if type(minQuantityRainfall) is np.str_ else minQuantityRainfall
        self.maxQR = maxQuantityRainfall.astype('float64') if type(maxQuantityRainfall) is np.str_ else maxQuantityRainfall
        self.sewageG = self.wc * self.popu * 0.85 * 0.001
        self.treatPlant = []
        self.o = [100*self.rc*self.da, 250*self.rc*self.da, 500*self.rc*self.da, 10*self.rc*self.da*self.minQR, 10*self.rc*self.da*self.maxQR]

    def addTreatmentPlant(self, factors:list) -> None:
        self.treatPlant.append(factors)
    
    def evaluate(self) -> None:
        self.pipedSewVolume = 0 # Volume of piped sewage
        allSewInput =  0
        allPlantDesignScale = 0
        for plant in self.treatPlant:
            self.pipedSewVolume += plant[1]*plant[2]/325
            allSewInput += plant[1]
            allPlantDesignScale += plant[0]
        self.dirSewVolume = self.sewageG - self.pipedSewVolume  # Volume of direct sewage discharge
        self.sewageCR = self.pipedSewVolume / self.sewageG      # Sewage collection rate
        self.pipelineBR = (allSewInput - self.pipedSewVolume) / allSewInput # Pipeline breakage rate
        self.j = [self.sewageG+i+self.pipedSewVolume*0.12 for i in self.o]
        self.h = allPlantDesignScale+(self.im+1)*self.pipedSewVolume
        self.r = [i - self.h for i in self.j]

def evaluate(original_data):

    result = {}
    for data in original_data:
        regionData, plantData = data[0:8], data[8:]
        region = Region(*regionData)

        plantData = np.delete(plantData, np.where(plantData == '/'))
        assert (len(plantData) % 3 == 0)
        plantData = [i.astype('float64') if type(i) is np.str_ else i for i in plantData]
        
        for factors in [plantData[i:i+3] for i in range(0, len(plantData) - 1, 3)]:
            if not (np.isnan(factors[0])):
                region.addTreatmentPlant(factors)
        region.evaluate()
        result[region.name] = {
            'Sewage_Generation' : region.sewageG,
            'Piped_Sewage_Volume' : region.pipedSewVolume,
            'Direct_Sewage_Discharge_Volume' : region.dirSewVolume,
            'Sewage_Collection_Rate' : region.sewageCR,
            'Pipeline_Breakage_Rate' : region.pipelineBR,
            'o' : region.o,
            'j' : region.j,
            'h' : region.h,
            'r' : region.r,
        }

    return result

if __name__ == "__main__":
    excel_path = "drainage-system/drainage_backend/evaluation/data.xlsx"
    df = pd.read_excel(io=excel_path)
    original_data = df.to_numpy()

    with open('./1.json', 'w') as file:
        json.dump(evaluate(original_data), file)
    pass